#include <cstring>
#include <functional>
#include <iostream>
#include <memory>
#include <random>
#include <set>
#include <string>

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")
// 为Windows平台定义ssize_t类型
#ifndef _SSIZE_T_DEFINED
typedef int ssize_t;
#define _SSIZE_T_DEFINED
#endif

#else
#include <arpa/inet.h>
#include <sys/socket.h>
#include <unistd.h>
#endif

// mDNS配置
constexpr char MDNS_GROUP[] = "224.0.0.251";
constexpr int MDNS_PORT = 5353;
#ifdef _WIN32
constexpr char LOCAL_INTERFACE[] = "192.168.22.178";// 替换为你的本地IP
#else
constexpr char LOCAL_INTERFACE[] = "192.168.22.13";// 替换为你的本地IP
#endif

// 创建mDNS PTR查询数据包
std::string create_mdns_query(const std::string &service_type)
{
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<uint16_t> dis(0, 0xFFFF);
    uint16_t query_id = dis(gen);

    uint16_t header[6] = {
            htons(query_id), htons(0x0100), htons(1), htons(0), htons(0), htons(0)};

    std::string qname;
    std::string remaining = service_type;
    size_t pos = 0;

    while ((pos = remaining.find('.', pos)) != std::string::npos)
    {
        qname += static_cast<char>(pos);
        qname += remaining.substr(0, pos);
        remaining = remaining.substr(pos + 1);
        pos = 0;
    }

    qname += static_cast<char>(remaining.length());
    qname += remaining;
    qname += '\x00';

    uint16_t qtype = htons(12);
    uint16_t qclass = htons(1);

    std::string question = qname;
    question.append(reinterpret_cast<const char *>(&qtype), sizeof(qtype));
    question.append(reinterpret_cast<const char *>(&qclass), sizeof(qclass));

    std::string query;
    query.append(reinterpret_cast<const char *>(header), sizeof(header));
    query += question;

    return query;
}

// 解析DNS名称（支持压缩指针）
std::pair<std::string, size_t> parse_dns_name(const std::string &data, size_t offset, const std::string &full_data)
{
    std::string labels;
    std::set<size_t> visited_pointers;

    while (offset < data.length())
    {
        uint8_t length = static_cast<uint8_t>(data[offset]);

        if (length == 0)
        {
            offset++;
            break;
        }
        else if ((length & 0xC0) == 0xC0)
        {
            if (offset + 1 >= data.length())
            {
                throw std::out_of_range("指针超出数据包范围");
            }

            size_t pointer = ((length & 0x3F) << 8) + static_cast<uint8_t>(data[offset + 1]);
            offset += 2;

            if (visited_pointers.find(pointer) != visited_pointers.end())
            {
                throw std::runtime_error("检测到指针循环引用");
            }

            visited_pointers.insert(pointer);

            if (pointer >= full_data.length())
            {
                throw std::out_of_range("指针指向完整数据包外");
            }

            std::pair<std::string, size_t> ptr_result = parse_dns_name(full_data, pointer, full_data);
            if (!labels.empty())
            {
                labels += ".";
            }
            labels += ptr_result.first;
            break;
        }
        else
        {
            uint8_t label_length = length;
            offset++;

            if (offset + label_length > data.length())
            {
                throw std::out_of_range("标签超出数据包范围");
            }

            if (!labels.empty())
            {
                labels += ".";
            }
            labels += data.substr(offset, label_length);
            offset += label_length;
        }
    }

    return {labels, offset};
}

// 解析mDNS响应数据包
void parse_mdns_response(const std::string &data)
{
    if (data.length() < 12)
    {
        std::cout << "[-] 无效的响应数据" << std::endl;
        return;
    }

    uint16_t header[6];
    memcpy(header, data.data(), sizeof(header));
    for (int i = 0; i < 6; ++i)
    {
        header[i] = ntohs(header[i]);
    }

    uint16_t qdcount = header[2];
    uint16_t ancount = header[3];
    uint16_t nscount = header[4];
    uint16_t arcount = header[5];

    size_t offset = 12;

    std::cout << "\n=== mDNS响应解析 ===" << std::endl;
    std::cout << "问题数量: " << qdcount << ", 回答数量: " << ancount << ", 权威记录数量: " << nscount << ", 附加记录数量: " << arcount << std::endl;

    // 跳过问题部分
    if (qdcount > 0)
    {
        std::cout << "\n--- 问题部分 ---" << std::endl;
        for (int i = 0; i < qdcount; ++i)
        {
            try
            {
                std::pair<std::string, size_t> result = parse_dns_name(data, offset, data);
                std::string name = result.first;
                offset = result.second;

                uint16_t qtype, qclass;
                memcpy(&qtype, data.data() + offset, sizeof(qtype));
                memcpy(&qclass, data.data() + offset + sizeof(qtype), sizeof(qclass));
                qtype = ntohs(qtype);
                qclass = ntohs(qclass);
                offset += 4;

                std::cout << "[" << i + 1 << "] " << name << ", QTYPE=" << qtype << ", QCLASS=" << qclass << std::endl;
            }
            catch (const std::exception &e)
            {
                std::cout << "[-] 解析问题 " << i + 1 << " 失败: " << e.what() << std::endl;
            }
        }
    }

    // 解析回答记录
    std::cout << "\n--- 回答记录 ---" << std::endl;
    for (int i = 0; i < ancount; ++i)
    {
        try
        {
            std::pair<std::string, size_t> result = parse_dns_name(data, offset, data);
            std::string name = result.first;
            offset = result.second;

            if (offset + 10 > data.length())
            {
                throw std::out_of_range("记录头部超出数据包范围");
            }

            uint16_t rtype, rclass;
            uint32_t ttl;
            uint16_t rdlength;
            memcpy(&rtype, data.data() + offset, sizeof(rtype));
            memcpy(&rclass, data.data() + offset + sizeof(rtype), sizeof(rclass));
            memcpy(&ttl, data.data() + offset + sizeof(rtype) + sizeof(rclass), sizeof(ttl));
            memcpy(&rdlength, data.data() + offset + sizeof(rtype) + sizeof(rclass) + sizeof(ttl), sizeof(rdlength));
            rtype = ntohs(rtype);
            rclass = ntohs(rclass);
            ttl = ntohl(ttl);
            rdlength = ntohs(rdlength);
            offset += 10;

            if (offset + rdlength > data.length())
            {
                throw std::out_of_range("记录数据超出数据包范围");
            }

            std::string rdata = data.substr(offset, rdlength);
            offset += rdlength;

            std::string record_info = "[" + std::to_string(i + 1) + "] " + name + ", TYPE=" + std::to_string(rtype) + ", CLASS=" + std::to_string(rclass) + ", TTL=" + std::to_string(ttl);

            if (rtype == 12)
            {// PTR记录
                std::pair<std::string, size_t> ptr_result = parse_dns_name(rdata, 0, data);
                std::string ptr_name = ptr_result.first;
                record_info += ", PTR=" + ptr_name;
            }
            else if (rtype == 1)
            {// A记录
                char ip[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, rdata.data(), ip, INET_ADDRSTRLEN);
                record_info += ", A=" + std::string(ip);
            }
            else if (rtype == 28)
            {// AAAA记录
                char ip[INET6_ADDRSTRLEN];
                inet_ntop(AF_INET6, rdata.data(), ip, INET6_ADDRSTRLEN);
                record_info += ", AAAA=" + std::string(ip);
            }
            else if (rtype == 33)
            {// SRV记录
                if (rdata.length() < 6)
                {
                    record_info += ", SRV=格式错误";
                }
                else
                {
                    uint16_t priority, weight, port;
                    memcpy(&priority, rdata.data(), sizeof(priority));
                    memcpy(&weight, rdata.data() + sizeof(priority), sizeof(weight));
                    memcpy(&port, rdata.data() + sizeof(priority) + sizeof(weight), sizeof(port));
                    priority = ntohs(priority);
                    weight = ntohs(weight);
                    port = ntohs(port);
                    std::pair<std::string, size_t> target_result = parse_dns_name(rdata, 6, data);
                    std::string target = target_result.first;
                    record_info += ", SRV=" + std::to_string(priority) + " " + std::to_string(weight) + " " + std::to_string(port) + " " + target;
                }
            }
            else if (rtype == 16)
            {// TXT记录
                std::string txt_data;
                size_t txt_offset = 0;
                while (txt_offset < rdata.length())
                {
                    uint8_t txt_len = static_cast<uint8_t>(rdata[txt_offset]);
                    txt_offset++;
                    if (txt_offset + txt_len > rdata.length())
                    {
                        break;
                    }
                    std::string txt_str = rdata.substr(txt_offset, txt_len);
                    txt_data += txt_str;
                    txt_offset += txt_len;
                    if (txt_offset < rdata.length())
                    {
                        txt_data += ", ";
                    }
                }
                record_info += ", TXT=[" + txt_data + "]";
            }

            std::cout << record_info << std::endl;
        }
        catch (const std::exception &e)
        {
            std::cout << "[-] 解析回答记录 " << i + 1 << " 失败: " << e.what() << std::endl;
        }
    }

    // 解析附加记录
    std::cout << "\n--- 附加记录 ---" << std::endl;
    for (int i = 0; i < arcount; ++i)
    {
        try
        {
            std::pair<std::string, size_t> result = parse_dns_name(data, offset, data);
            std::string name = result.first;
            offset = result.second;

            if (offset + 10 > data.length())
            {
                throw std::out_of_range("附加记录头部超出数据包范围");
            }

            uint16_t rtype, rclass;
            uint32_t ttl;
            uint16_t rdlength;
            memcpy(&rtype, data.data() + offset, sizeof(rtype));
            memcpy(&rclass, data.data() + offset + sizeof(rtype), sizeof(rclass));
            memcpy(&ttl, data.data() + offset + sizeof(rtype) + sizeof(rclass), sizeof(ttl));
            memcpy(&rdlength, data.data() + offset + sizeof(rtype) + sizeof(rclass) + sizeof(ttl), sizeof(rdlength));
            rtype = ntohs(rtype);
            rclass = ntohs(rclass);
            ttl = ntohl(ttl);
            rdlength = ntohs(rdlength);
            offset += 10;

            if (offset + rdlength > data.length())
            {
                throw std::out_of_range("附加记录数据超出数据包范围");
            }

            std::string rdata = data.substr(offset, rdlength);
            offset += rdlength;

            std::string record_info = "[" + std::to_string(i + 1) + "] " + name + ", TYPE=" + std::to_string(rtype) + ", CLASS=" + std::to_string(rclass) + ", TTL=" + std::to_string(ttl);

            if (rtype == 1)
            {// A记录
                char ip[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, rdata.data(), ip, INET_ADDRSTRLEN);
                record_info += ", A=" + std::string(ip);
            }
            else if (rtype == 28)
            {// AAAA记录
                char ip[INET6_ADDRSTRLEN];
                inet_ntop(AF_INET6, rdata.data(), ip, INET6_ADDRSTRLEN);
                record_info += ", AAAA=" + std::string(ip);
            }
            else if (rtype == 33)
            {// SRV记录
                if (rdata.length() < 6)
                {
                    record_info += ", SRV=格式错误";
                }
                else
                {
                    uint16_t priority, weight, port;
                    memcpy(&priority, rdata.data(), sizeof(priority));
                    memcpy(&weight, rdata.data() + sizeof(priority), sizeof(weight));
                    memcpy(&port, rdata.data() + sizeof(priority) + sizeof(weight), sizeof(port));
                    priority = ntohs(priority);
                    weight = ntohs(weight);
                    port = ntohs(port);
                    std::pair<std::string, size_t> target_result = parse_dns_name(rdata, 6, data);
                    std::string target = target_result.first;
                    record_info += ", SRV=" + std::to_string(priority) + " " + std::to_string(weight) + " " + std::to_string(port) + " " + target;
                }
            }
            else if (rtype == 16)
            {// TXT记录
                std::string txt_data;
                size_t txt_offset = 0;
                while (txt_offset < rdata.length())
                {
                    uint8_t txt_len = static_cast<uint8_t>(rdata[txt_offset]);
                    txt_offset++;
                    if (txt_offset + txt_len > rdata.length())
                    {
                        break;
                    }
                    std::string txt_str = rdata.substr(txt_offset, txt_len);
                    txt_data += txt_str;
                    txt_offset += txt_len;
                    if (txt_offset < rdata.length())
                    {
                        txt_data += ", ";
                    }
                }
                record_info += ", TXT=[" + txt_data + "]";
            }

            std::cout << record_info << std::endl;
        }
        catch (const std::exception &e)
        {
            std::cout << "[-] 解析附加记录 " << i + 1 << " 失败: " << e.what() << std::endl;
        }
    }
}

// 资源管理类，用于自动关闭socket
class Socket
{
private:
    int sock;

public:
    Socket(int domain, int type, int protocol) : sock(socket(domain, type, protocol))
    {
        if (sock < 0)
        {
            throw std::runtime_error("Failed to create socket");
        }
    }

    ~Socket()
    {
#ifdef _WIN32
        if (sock >= 0)
        {
            closesocket(sock);
        }
#else
        if (sock >= 0)
        {
            close(sock);
        }
#endif
    }

    // 禁用拷贝
    Socket(const Socket &) = delete;
    Socket &operator=(const Socket &) = delete;

    // 允许移动
    Socket(Socket &&other) noexcept : sock(other.sock)
    {
        other.sock = -1;
    }

    Socket &operator=(Socket &&other) noexcept
    {
        if (this != &other)
        {
#ifdef _WIN32
            if (sock >= 0)
            {
                closesocket(sock);
            }
#else
            if (sock >= 0)
            {
                close(sock);
            }
#endif
            sock = other.sock;
            other.sock = -1;
        }
        return *this;
    }

    int get() const { return sock; }
};

// 发送mDNS查询并接收响应
std::string send_and_receive_mdns_query(const std::string &service_type)
{
    try
    {
#ifdef _WIN32
        WSADATA wsaData;
        if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
        {
            throw std::runtime_error("Failed to initialize Winsock");
        }
#endif

        // 创建发送套接字
        Socket send_sock(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

        int reuse = 1;
// 修复：在Windows平台显式转换为const char*
#ifdef _WIN32
        if (setsockopt(send_sock.get(), SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<const char *>(&reuse), sizeof(reuse)) < 0)
#else
        if (setsockopt(send_sock.get(), SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0)
#endif
        {
            throw std::runtime_error("Failed to set socket option");
        }

        struct in_addr localInterface;
        localInterface.s_addr = inet_addr(LOCAL_INTERFACE);
        if (setsockopt(send_sock.get(), IPPROTO_IP, IP_MULTICAST_IF, reinterpret_cast<char *>(&localInterface), sizeof(localInterface)) < 0)
        {
            throw std::runtime_error("Failed to set multicast interface");
        }

        // 发送查询
        std::string query = create_mdns_query(service_type);
        struct sockaddr_in groupSock;
        memset(&groupSock, 0, sizeof(groupSock));
        groupSock.sin_family = AF_INET;
        groupSock.sin_addr.s_addr = inet_addr(MDNS_GROUP);
        groupSock.sin_port = htons(MDNS_PORT);

        if (sendto(send_sock.get(), query.c_str(), query.length(), 0, reinterpret_cast<struct sockaddr *>(&groupSock), sizeof(groupSock)) < 0)
        {
            throw std::runtime_error("Failed to send mDNS query");
        }
        std::cout << "[+] 发送mDNS查询: " << service_type << std::endl;

        // 创建接收套接字
        Socket recv_sock(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

// 修复：在Windows平台显式转换为const char*
#ifdef _WIN32
        if (setsockopt(recv_sock.get(), SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<const char *>(&reuse), sizeof(reuse)) < 0)
#else
        if (setsockopt(recv_sock.get(), SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0)
#endif
        {
            throw std::runtime_error("Failed to set socket option");
        }

        struct sockaddr_in localSock;
        memset(&localSock, 0, sizeof(localSock));
        localSock.sin_family = AF_INET;
        localSock.sin_port = htons(MDNS_PORT);
        localSock.sin_addr.s_addr = INADDR_ANY;

        if (bind(recv_sock.get(), reinterpret_cast<struct sockaddr *>(&localSock), sizeof(localSock)) < 0)
        {
            throw std::runtime_error("Failed to bind socket");
        }

        // 加入多播组
        struct ip_mreq group;
        group.imr_multiaddr.s_addr = inet_addr(MDNS_GROUP);
        group.imr_interface.s_addr = INADDR_ANY;
        if (setsockopt(recv_sock.get(), IPPROTO_IP, IP_ADD_MEMBERSHIP, reinterpret_cast<char *>(&group), sizeof(group)) < 0)
        {
            throw std::runtime_error("Failed to join multicast group");
        }

        // 设置超时时间
        struct timeval timeout;
        timeout.tv_sec = 10;// 延长超时时间到10秒
        timeout.tv_usec = 0;

// 修复：在Windows平台显式转换为const char*
#ifdef _WIN32
        if (setsockopt(recv_sock.get(), SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<const char *>(&timeout), sizeof(timeout)) < 0)
#else
        if (setsockopt(recv_sock.get(), SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<char *>(&timeout), sizeof(timeout)) < 0)
#endif
        {
            throw std::runtime_error("Failed to set receive timeout");
        }

        char buffer[4096];
        struct sockaddr_in sender;
        socklen_t senderLen = sizeof(sender);
        ssize_t recvLen = recvfrom(recv_sock.get(), buffer, sizeof(buffer), 0, reinterpret_cast<struct sockaddr *>(&sender), &senderLen);
        if (recvLen < 0)
        {
#ifdef _WIN32
            if (WSAGetLastError() == WSAETIMEDOUT)
            {
                std::cout << "[-] 查询超时，未收到响应" << std::endl;
            }
            else
            {
                throw std::runtime_error("Failed to receive mDNS response");
            }
#else
            if (errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::cout << "[-] 查询超时，未收到响应" << std::endl;
            }
            else
            {
                throw std::runtime_error("Failed to receive mDNS response");
            }
#endif
            return "";
        }
        else
        {
            std::cout << "[+] 收到来自 " << inet_ntoa(sender.sin_addr) << ":" << ntohs(sender.sin_port) << " 的响应" << std::endl;
            std::string hex_data;
            for (ssize_t i = 0; i < recvLen; ++i)
            {
                char hex[3];
                sprintf(hex, "%02x", static_cast<unsigned char>(buffer[i]));
                hex_data += hex;
            }
            std::cout << "\n原始数据:\n"
                      << hex_data << std::endl;
            return std::string(buffer, recvLen);
        }

#ifdef _WIN32
        WSACleanup();
#endif
    }
    catch (const std::exception &e)
    {
        std::cerr << "[-] 错误: " << e.what() << std::endl;
#ifdef _WIN32
        WSACleanup();
#endif
        return "";
    }
}

int mdns_query_parse()
{
    try
    {
        std::string service_type = "_smartaudio._udp.local";
        std::string response = send_and_receive_mdns_query(service_type);

        if (!response.empty())
        {
            parse_mdns_response(response);
        }
        else
        {
            std::cout << "[-] 未获取到有效响应" << std::endl;
        }

        return 0;
    }
    catch (const std::exception &e)
    {
        std::cerr << "[-] 程序异常终止: " << e.what() << std::endl;
        return 1;
    }
}
